'''
@Project ：conda_envs
@File    ：dataProcessing.py
@IDE     ：PyCharm
@Author  ：子协
@Date    ：2025-02-27 21:10
'''
# 对脑电原始信号进行预处理，提取出特征
# 流程：读取原始数据 -> 带通滤波去噪 -> [主成分分析去伪影] -> 特征提取（时域，频域，时频域）

import numpy as np
from scipy import signal, stats

class EEGProcessor:
    def __init__(self):
        self.rawData = np.zeros(0)                  # 原始数据
        self.baseline = np.zeros(0)                 # 基线数据
        self.baselineCorrection_rawData = np.zeros(0)# 基线矫正后的原始数据
        self.filter_rawData = np.zeros(0)           # 滤波后的原始数据
        self.filter_baseline = np.zeros(0)          # 滤波后的基线数据
        self.timedomainProperty = np.zeros(0)       # 时域特征
        self.fredomainProperty = np.zeros(0)        # 频域特征
        self.property = np.zeros(0)                 # 所有特征（时域，频域，时频域）
        self.scaleProperty = np.zeros(0)            # 归一化后的数据

    def loadData(self,rawData,baseLine):
        self.rawData = np.array(rawData)
        self.baseline = np.array(baseLine)

    def clearData(self):
        np.delete(self.rawData, np.arange(len(self.rawData)))
        np.delete(self.filter_rawData, np.arange(len(self.filter_rawData)))
        np.delete(self.timedomainProperty, np.arange(len(self.timedomainProperty)))
        np.delete(self.fredomainProperty, np.arange(len(self.fredomainProperty)))
        np.delete(self.scaleProperty, np.arange(len(self.scaleProperty)))
        np.delete(self.property, np.arange(len(self.property)))
        np.delete(self.baselineCorrection_rawData,np.arange(len(self.baselineCorrection_rawData)))
        np.delete(self.baseline, np.arange(len(self.baseline)))
        np.delete(self.filter_baseline, np.arange(len(self.filter_baseline)))

    ''' 
    * @brief          获取基线的均值
    * @param baseline 基线数据(二维数组)
    * @return         基线的均值
    '''
    def GetBaseLineMean(self,baseline):
        _baseline = np.array(baseline)
        _baseline = _baseline.flatten()  # 展平为一维
        self.baseline = _baseline        # 获取基线数据
        mean = np.around(_baseline.mean(), 6)
        return mean
    ''' 
    * @brief          基线矫正
    * @param 
    * @return         矫正后的数据
    '''
    def BaselineCorrection(self):
        mean = np.around(self.baseline.mean(), 6)
        self.baselineCorrection_rawData = self.rawData - mean
        _baselineCorrection_rawData = self.baselineCorrection_rawData.tolist()
        return _baselineCorrection_rawData

    ''' 
    * @brief          带通滤波
    * @param lowcut   频率下限
    * @param highcut  频率上线
    * @param fs       采样频率
    * @param order    阶数
    * @return         去噪后的数据
    '''
    def bandpassFilter(self, lowcut = 4, highcut = 30, fs=512, order=4):
        nyquist = 0.5 * fs
        low = lowcut / nyquist
        high = highcut / nyquist

        b, a = signal.butter(N = order, Wn = [low, high], btype='bandpass')

        self.filter_rawData = signal.filtfilt(b, a, self.baselineCorrection_rawData)
        self.filter_baseline = signal.filtfilt(b, a, self.baseline)

        _filter_rawData = self.filter_rawData.tolist()
        _filter_baseline = self.filter_baseline.tolist()
        return _filter_rawData,_filter_baseline

    ''' 
    * @brief        提取时域特征
    * @param 
    * @return       均值 方差 峰度 偏度 均方根
    '''
    def TimeDomain(self):
         _mean = np.mean(self.filter_rawData)               # 均值
         _var = np.var(self.filter_rawData)                 # 方差
         _kurtosis = stats.kurtosis(self.filter_rawData)    # 峰度
         _skew = stats.skew(np.array(self.filter_rawData))  # 偏度
         # _sqrt = np.sqrt(self.filter_rawData)               # 均方根
         _timedomainProperty = [_mean, _var, _kurtosis, _skew]
         self.timedomainProperty = np.array(_timedomainProperty)
         return self.timedomainProperty

    ''' 
    * @brief        提取频域特征
    * @param   
    * @return       各频段的能量(绝对、相对)，功率谱密度和功率谱熵
    '''
    def FreDomain(self):
        bands = {
            "theta":(4,8),
            "alpha":(8, 13),
            "beta":(13, 30),
            "gamma":(30, 45)
        }
        _fredomainProperty = []
        # 计算功率谱密度（Welch方法）
        freqs_rawData, psd_rawData = signal.welch(self.filter_rawData, fs=512, nperseg=256)        # 滤波后的原始数据功率谱密度
        freqs_baseLine, psd_baseLine = signal.welch(self.filter_baseline, fs=512, nperseg=256)     # 滤波后的基线数据功率谱密度
        theta_power_after = np.sum(psd_rawData[(freqs_rawData >= bands["theta"][0]) & (freqs_rawData <= bands["theta"][1])])     # 事件后θ波能量
        alpha_power_after = np.sum(psd_rawData[(freqs_rawData >= bands["alpha"][0]) & (freqs_rawData <= bands["alpha"][1])])     # 事件后α波能量
        beta_power_after = np.sum(psd_rawData[(freqs_rawData >= bands["beta"][0]) & (freqs_rawData <= bands["beta"][1])])        # 事件后β波能量
        gamma_power_after = np.sum(psd_rawData[(freqs_rawData >= bands["gamma"][0]) & (freqs_rawData <= bands["gamma"][1])])     # 事件后γ波能量
        alpha_power_before = np.sum(psd_baseLine[(freqs_baseLine >= bands["alpha"][0]) & (freqs_baseLine <= bands["alpha"][1])])  # 事件前α波能量
        beta_power_before = np.sum(psd_baseLine[(freqs_baseLine >= bands["beta"][0]) & (freqs_baseLine <= bands["beta"][1])])   # 事件前β波能量

        beta_alpha_ratio = beta_power_after / (alpha_power_after + 1e-10)                                      # 事件后β/α频段能量比
        beta_theta_ratio = beta_power_after / (theta_power_after + 1e-10)                                      # 事件后β/θ频段能量比

        alpha_avg_power_after = np.mean(psd_rawData[(freqs_rawData >= bands["alpha"][0]) & (freqs_rawData <= bands["alpha"][1])])# 事件后α波平均功率
        alpha_max_power_after = np.max(psd_rawData[(freqs_rawData >= bands["alpha"][0]) & (freqs_rawData <= bands["alpha"][1])])  # 事件后α波最大功率
        beta_avg_power_after = np.mean(psd_rawData[(freqs_rawData >= bands["beta"][0]) & (freqs_rawData <= bands["beta"][1])])    # 事件后β波平均功率
        beta_max_power_after = np.max(psd_rawData[(freqs_rawData >= bands["beta"][0]) & (freqs_rawData <= bands["beta"][1])])     # 事件后β波最大功率

        ERS_ERD_alpha = (alpha_power_after - alpha_power_before) / alpha_power_before             # ERS/ERD 事件发生后的α波能量/事件发生前α波能量
        ERS_ERD_beta = (beta_power_after - beta_power_before) / beta_power_before               # ERS/ERD 事件发生后的β波能量/事件发生前α波能量

        # _fredomainProperty.extend(psd_rawData)
        # _fredomainProperty.append(theta_power_after)
        _fredomainProperty.append(alpha_power_after)
        _fredomainProperty.append(beta_power_after)
        # _fredomainProperty.append(gamma_power_after)
        _fredomainProperty.append(beta_alpha_ratio)
        # _fredomainProperty.append(beta_theta_ratio)
        _fredomainProperty.append(alpha_avg_power_after)
        _fredomainProperty.append(alpha_max_power_after)
        _fredomainProperty.append(beta_avg_power_after)
        _fredomainProperty.append(beta_max_power_after)
        # _fredomainProperty.append(ERS_ERD_alpha)
        # _fredomainProperty.append(ERS_ERD_beta)

        self.fredomainProperty = np.array(_fredomainProperty)
        return _fredomainProperty

    ''' 
    * @brief        提取时频特征
    * @param   
    * @return       
    '''
    def TimeFreDomain(self):
        pass

    ''' 
    * @brief        获取所有特征（时域，频域，时频域）
    * @param   
    * @return       所有特征（时域，频域，时频域）
    '''
    def GetProperty(self):
        self.bandpassFilter()
        self.TimeDomain()
        self.FreDomain()
        self.property = np.concatenate((self.timedomainProperty,self.fredomainProperty))
        _property = self.property.tolist()
        return _property

    ''' 
    * @brief        数据归一化(模型训练时)
    * @param 
    * @return       归一化后的数据和最大值最小值
    '''
    def scale_for_train(self, data, range=(-1, 1)):
        X = np.array(data)
        upper = range[1]
        lower = range[0]
        feat_min = X.min(axis=0)  # 计算每列（特征）的最小值和最大值
        feat_max = X.max(axis=0)
        scale = (upper - lower) / (feat_max - feat_min)  # 计算缩放系数
        offset = lower - feat_min * scale  # 计算偏移量
        X_scaled = X * scale + offset  # 对数据进行缩放处理
        X_scaled = np.round(X_scaled, 6)    # 保留6位小数
        feat_min = np.round(feat_min, 6)    # 保留6位小数
        feat_max = np.round(feat_max, 6)    # 保留6位小数
        X_scaled = X_scaled.tolist()
        _feat_min = feat_min.tolist()
        _feat_max = feat_max.tolist()
        return X_scaled,_feat_min,_feat_max

    def scale_for_predict(self, data, feat_min,feat_max,range=(-1, 1)):
        X = np.array(data)
        upper = range[1]
        lower = range[0]
        _feat_min = np.array(feat_min)
        _feat_max = np.array(feat_max)
        scale = (upper - lower) / (_feat_max - _feat_min)  # 计算缩放系数
        offset = lower - _feat_min * scale  # 计算偏移量
        X_scaled = X * scale + offset  # 对数据进行缩放处理
        X_scaled = X_scaled.tolist()
        return X_scaled

